import json

seg_result = []

def clamp(x, s, e):
    if(x <= s):
        x = s
    if(x >= e):
        x = e
    return x

def get_first_seg(name, height, width, category, p1, p2):
    x1, y1 = p1[0], p1[1]
    x2, y2 = p2[0], p2[1]
    h = height // 2
    w = width // 2
    if(x1 < w and y1 < h):
        result = {}
        name_now = name.split('.')[0] + '_1.' + name.split('.')[1]
        result['name'] = name_now
        result['image_height'] = h
        result['image_width'] = w
        result['category'] = category
        bbox = []
        point1 = (x1, y1)
        bbox.append(point1[0])
        bbox.append(point1[1])
        point2 = []
        if (x2 < w):
            point2.append(x2)
        else:
            point2.append(w)
        if (y2 < h):
            point2.append(y2)
        else:
            point2.append(h)
        bbox.append(point2[0])
        bbox.append(point2[1])
        for i, x in enumerate(bbox):
            if(i == 0 or i == 2):
                bbox[i] = clamp(x, 0, w)
            else:
                bbox[i] = clamp(x, 0, h)
        result['bbox'] = bbox
        assert bbox[0] <= w and bbox[1] <= h and bbox[2] <= w and bbox[3] <= h
        seg_result.append(result)

def get_second_seg(name, height, width, category, p1, p2):
    x1, y1 = p1[0], p1[1]
    x2, y2 = p2[0], p2[1]
    h = height // 2
    w = width // 2
    if(x2 > w and y1 < h):
        result = {}
        name_now = name.split('.')[0] + '_2.' + name.split('.')[1]
        result['name'] = name_now
        result['image_height'] = h
        result['image_width'] = w
        result['category'] = category
        bbox = []
        point1 = []
        point2 = []
        if (x1 > w):
            point1.append(x1 - w)
            point1.append(y1)
        else:
            point1.append(0)
            point1.append(y1)
        if (y2 < h):
            point2.append(x2 - w)
            point2.append(y2)
        else:
            point2.append(x2 - w)
            point2.append(h)
        for i in point1:
            bbox.append(i)
        for i in point2:
            bbox.append(i)
        for i, x in enumerate(bbox):
            if(i == 0 or i == 2):
                bbox[i] = clamp(x, 0, w)
            else:
                bbox[i] = clamp(x, 0, h)
        result['bbox'] = bbox
        assert bbox[0] <= w and bbox[1] <= h and bbox[2] <= w and bbox[3] <= h
        seg_result.append(result)

def get_third_seg(name, height, width, category, p1, p2):
    x1, y1 = p1[0], p1[1]
    x2, y2 = p2[0], p2[1]
    h = height // 2
    w = width // 2
    if(x1 < w and y2 > h):
        result = {}
        name_now = name.split('.')[0] + '_3.' + name.split('.')[1]
        result['name'] = name_now
        result['image_height'] = h
        result['image_width'] = w
        result['category'] = category
        bbox = []
        point1 = []
        point2 = []
        if (y1 > h):
            point1.append(x1)
            point1.append(y1 - h)
        else:
            point1.append(x1)
            point1.append(0)
        if (x2 < w):
            point2.append(x2)
            point2.append(y2 - h)
        else:
            point2.append(w)
            point2.append(y2 - h)
        for i in point1:
            bbox.append(i)
        for i in point2:
            bbox.append(i)
        for i, x in enumerate(bbox):
            if(i == 0 or i == 2):
                bbox[i] = clamp(x, 0, w)
            else:
                bbox[i] = clamp(x, 0, h)
        result['bbox'] = bbox
        assert bbox[0] <= w and bbox[1] <= h and bbox[2] <= w and bbox[3] <= h
        seg_result.append(result)

def get_fourth_seg(name, height, width, category, p1, p2):
    x1, y1 = p1[0], p1[1]
    x2, y2 = p2[0], p2[1]
    h = height // 2
    w = width // 2
    if(x2 > w and y2 > h):
        result = {}
        name_now = name.split('.')[0] + '_4.' + name.split('.')[1]
        result['name'] = name_now
        result['image_height'] = h
        result['image_width'] = w
        result['category'] = category
        bbox = []
        point1 = []
        point2 = (x2 - w, y2 - h)
        if (x1 > w):
            point1.append(x1 - w)
        else:
            point1.append(0)
        if (y1 > h):
            point1.append(y1 - h)
        else:
            point1.append(0)
        for i in point1:
            bbox.append(i)
        for i in point2:
            bbox.append(i)
        for i, x in enumerate(bbox):
            if(i == 0 or i == 2):
                bbox[i] = clamp(x, 0, w)
            else:
                bbox[i] = clamp(x, 0, h)
        result['bbox'] = bbox
        assert bbox[0] <= w and bbox[1] <= h and bbox[2] <= w and bbox[3] <= h
        seg_result.append(result)

if __name__ == "__main__":
    # json_dir = '/data/lzy/tianchi2021/train_1to4_annos.json'
    json_dir = '/data/lzy/tile_round2_train_20210204/train_annos.json'
    # json_dir = '/data/lzy/tianchi2021/all_coco.json'
    save_dir = '/data/lzy/tile_round2_train_20210204/train_1to4_annos.json'
    with open(json_dir) as f:
        annos = json.load(f)
    for anno in annos:
        name = anno['name']
        image_height = anno['image_height']
        image_width = anno['image_width']
        category = anno['category']
        bbox = anno['bbox']
        p1 = (bbox[0], bbox[1])
        p2 = (bbox[2], bbox[3])
        if(image_width >= 1500 or image_height >= 1500):
            get_first_seg(name, image_height, image_width, category, p1, p2)
            get_second_seg(name, image_height, image_width, category, p1, p2)
            get_third_seg(name, image_height, image_width, category, p1, p2)
            get_fourth_seg(name, image_height, image_width, category, p1, p2)
    with open(save_dir, 'w') as fp:
        json.dump(seg_result, fp)